// gene expression set

#ifndef EXPRESSION_HH_
#define EXPRESSION_HH_
#include "util.hh"
#include <boost/shared_ptr.hpp>
#include <boost/unordered_map.hpp>

template<typename T>
struct ptr_hash_t
{
  size_t operator()(const typename boost::shared_ptr<T> obj_pt) const
  {
    return obj_pt->id;
  }
};

template<typename EXPR>
struct expr_ptr_set_t
{
private:

  typedef typename boost::unordered_map<boost::shared_ptr<EXPR>, double,
      ptr_hash_t<EXPR> > data_map_t;

  data_map_t data_map;

public:

  typedef boost::shared_ptr<EXPR> key_type;
  typedef typename data_map_t::const_iterator iterator_t;

  explicit expr_ptr_set_t() :
      data_map()
  {
  }

  void insert(key_type pt, double p = 1.)
  {
    data_map[pt] = p;
  }

  void remove(key_type pt)
  {
    data_map.erase(pt);
  }

  void clear()
  {
    data_map.clear();
  }

  size_t size() const
  {
    return data_map.size();
  }

  iterator_t begin() const
  {
    return data_map.cbegin();
  }

  iterator_t end() const
  {
    return data_map.cend();
  }

};

template<typename T>
struct expr_ptr_set_iterator_t
{
private:
  typedef typename expr_ptr_set_t<T>::iterator_t it_t;
  it_t it;
public:
  explicit expr_ptr_set_iterator_t(it_t _it) :
      it(_it)
  {
  }

  bool operator ==(const expr_ptr_set_iterator_t& rhs) const
  {
    return it == rhs.it;
  }

  bool operator !=(const expr_ptr_set_iterator_t& rhs) const
  {
    return !(*this == rhs);
  }

  expr_ptr_set_iterator_t& operator ++()
  {
    ++it;
    return *this;
  }

  const T& operator *() const
  {
    return *(it->first.get());
  }

  double get_prob()
  {
    return it->second;
  }
};

////////////////////////////////////////////////////////////////
// for a linear regression model
struct expr_lm_t
{
  typedef std::vector<double> vec_type;
  typedef double value_type;

  explicit expr_lm_t(const int _id, const size_t _p) :
      id(_id), p(_p), x(p, 0.), y(0.)
  {
  }

  const int id;
  const size_t p;
  vec_type x;
  value_type y;
};

typedef expr_ptr_set_t<expr_lm_t> expr_lm_set_t;
typedef expr_ptr_set_iterator_t<expr_lm_t> expr_lm_set_iter_t;

expr_lm_set_iter_t begin(const expr_lm_set_t& mat)
{
  return expr_lm_set_iter_t(mat.begin());
}

expr_lm_set_iter_t end(const expr_lm_set_t& mat)
{
  return expr_lm_set_iter_t(mat.end());
}

////////////////////////////////////////////////////////////////
// to represent expression matrix
struct expr_vector_t
{
  explicit expr_vector_t(const int _id, const size_t _p) :
      id(_id), p(_p), vec(_p, 0.)
  {
  }

  const int id;
  const size_t p;
  std::vector<double> vec;
};

typedef expr_ptr_set_t<expr_vector_t > expr_vec_set_t;
typedef expr_ptr_set_iterator_t<expr_vector_t > expr_vec_set_iter_t;

expr_vec_set_iter_t begin(const expr_vec_set_t& mat)
{
  return expr_vec_set_iter_t(mat.begin());
}

expr_vec_set_iter_t end(const expr_vec_set_t& mat)
{
  return expr_vec_set_iter_t(mat.end());
}

////////////////////////////////////////////////////////////////
// I/O routine
typedef std::vector<boost::shared_ptr<expr_lm_t> > expr_lm_vec_t;

bool read_expr_vec(const char* x_file,
    std::vector<boost::shared_ptr<expr_vector_t> >& out, size_t& n, size_t& p)
{
  p = num_columns(x_file);
  n = num_rows(x_file);

  TLOG("X [" << n << " x " << p << "]");

  out.clear();

  std::ifstream xstr(x_file, std::ios::in);

  double x;

  int id = -1;
  for (size_t i = 0; (i < n * p) && (xstr >> x); ++i)
  {
    size_t pos = i % p;

    if (pos == 0)
      out.push_back(boost::shared_ptr<expr_vector_t>(new expr_vector_t(++id, p)));

    expr_vector_t& data = *(out[id].get());
    data.vec[pos] = x;
  }

  xstr.close();

  if (id != (n - 1))
    return false;

  return true;
}

bool read_expr_lm_vec(const char* x_file, const char* y_file,
    expr_lm_vec_t& out, size_t& n, size_t& p)
{
  p = num_columns(x_file);
  n = num_rows(x_file);

  TLOG("X [" << n << " x " << p << "]");

  assert_msg(n == num_rows(y_file), "number of data points should match");

  out.clear();

  std::ifstream xstr(x_file, std::ios::in);
  std::ifstream ystr(y_file, std::ios::in);

  double x, y;

  int id = -1;
  for (size_t i = 0; (i < n * p) && (xstr >> x); ++i)
  {
    size_t pos = i % p;

    if (pos == 0)
    {
      // add intercept
      out.push_back(boost::shared_ptr<expr_lm_t>(new expr_lm_t(++id, p + 1)));

      if (!(ystr >> y))
      {
        xstr.close();
        ystr.close();
        return false;
      }

      expr_lm_t& lm_obj = *(out[id].get());
      lm_obj.y = y;
      lm_obj.x[pos] = 1.;
      lm_obj.x[pos + 1] = x;
    }
    else
    {
      expr_lm_t& lm_obj = *(out[id].get());
      lm_obj.x[pos + 1] = x;
    }
  }

  xstr.close();
  ystr.close();

  if (id != (n - 1))
    return false;

  ++p;  // add intercept variable
  return true;
}

#endif /* EXPRESSION_HH_ */
